#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
modbus_cve_probe_runner.py  (no-CLI, live-print)

- 不依赖任何外部输入参数；配置写死在文件顶部 CONFIG 中。
- 保持实时打印命中与 JSONL 证据输出；绝不修改你的原始服务代码，仅以 middleware 形式挂载。

覆盖 CVE（启发式签名）:
- CVE-2022-28613  输入验证不当（重启）
- CVE-2024-34244  缓冲区溢出（越界读）
- CVE-2024-10918  栈缓冲区溢出
- CVE-2023-25619  异常条件检查不当，DoS
"""

from __future__ import annotations

import asyncio
import json
from datetime import datetime, timezone
from typing import Callable, Awaitable, Dict, Any, Optional, Tuple
import struct
from collections import defaultdict

# --- 固定配置（不要外部输入参数） ---
CONFIG = {
    "HOST": "0.0.0.0",
    "PORT": None,              # 若为 None，则取被测模块的 MODBUS_TCP_DEFAULT_PORT 或 1502
    "OUT": "cve_hits.jsonl",
    "RATELIMIT": 200,          # requests per second
}

# --- 导入你的 Modbus 栈（不做任何修改） ---
try:
    import modbus_full_stack_schemeB as mbs  # type: ignore
except Exception:
    import modbus_full_stack as mbs  # type: ignore

MBAP = mbs.MBAP
RequestContext = mbs.RequestContext
REG_CNT_MAX = getattr(mbs, "REG_CNT_MAX", 2000)

def _resolve_port() -> int:
    if CONFIG["PORT"] is not None:
        return int(CONFIG["PORT"])
    return int(getattr(mbs, "MODBUS_TCP_DEFAULT_PORT", 1503))

def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat(timespec="milliseconds")

def ts_hms() -> str:
    return datetime.now().strftime("%H:%M:%S.%f")[:-3]

def hexdump(data: bytes) -> str:
    return " ".join(f"{b:02X}" for b in data)

def safe_unpack(fmt: str, buf: bytes) -> Optional[Tuple]:
    try:
        size = struct.calcsize(fmt)
        if len(buf) < size:
            return None
        return struct.unpack(fmt, buf[:size])
    except Exception:
        return None

def pdu_tail(pdu: bytes, offset: int) -> bytes:
    return pdu[offset:] if len(pdu) >= offset else b""

# ----------------------------------
# CVE Detection
# ----------------------------------

class CVELog:
    def __init__(self, path: str) -> None:
        self.path = path
        self._fh = open(self.path, "a", encoding="utf-8")

    def write(self, entry: Dict[str, Any]) -> None:
        self._fh.write(json.dumps(entry, ensure_ascii=False) + "\n")
        self._fh.flush()

    def close(self) -> None:
        try:
            self._fh.close()
        except Exception:
            pass

class CVEProbe:
    """
    Stateless signature checks that run on every request (in middleware).
    Adds live printing on each hit and keeps aggregate counters.
    """
    def __init__(self, logger: CVELog) -> None:
        self.log = logger
        self.counts = defaultdict(int)  # cve -> count

    def _emit(self, ctx: RequestContext, cve: str, title: str, reason: str) -> None:
        mbap: MBAP = ctx.mbap  # type: ignore[assignment]
        entry = {
            "ts": now_iso(),
            "peer": ctx.peername,
            "cve": cve,
            "title": title,
            "reason": reason,
            "tid": int(getattr(mbap, "tid", -1)),
            "pid": int(getattr(mbap, "pid", -1)),
            "len": int(getattr(mbap, "length", -1)),
            "uid": int(getattr(mbap, "uid", -1)),
            "fc": int(ctx.fc),
            "pdu_hex": hexdump(ctx.pdu),
        }
        # persist
        self.log.write(entry)
        # aggregate
        self.counts[cve] += 1
        # live print
        totals = " | ".join(f"{k}:{self.counts[k]}" for k in sorted(self.counts))
        print(f"[{ts_hms()}] [CVE HIT] {cve} {title} | fc=0x{ctx.fc:02X} | {reason} | totals: {totals}")

    async def check(self, ctx: RequestContext) -> None:
        self._check_mbap(ctx)
        fc = ctx.fc
        if fc in (0x01, 0x02, 0x03, 0x04):
            self._check_read_blocks(ctx, fc)
        elif fc == 0x05:
            self._check_write_single_coil(ctx)
        elif fc == 0x06:
            self._check_write_single_reg(ctx)
        elif fc == 0x0F:
            self._check_write_multi_coils(ctx)
        elif fc == 0x10:
            self._check_write_multi_regs(ctx)
        elif fc == 0x16:
            self._check_mask_write_reg(ctx)
        elif fc == 0x17:
            self._check_read_write_multi(ctx)

    # --- Individual detectors ---
    def _check_mbap(self, ctx: RequestContext) -> None:
        mbap: MBAP = ctx.mbap  # type: ignore[assignment]
        if mbap.pid != 0:
            self._emit(ctx, "CVE-2022-28613", "输入验证不当（重启）", f"MBAP.pid={mbap.pid} (should be 0)")
        if mbap.length == 0:
            self._emit(ctx, "CVE-2022-28613", "输入验证不当（重启）", "MBAP.length=0")
        if mbap.length > 1024:
            self._emit(ctx, "CVE-2023-25619", "异常条件检查不当，DoS", f"MBAP.length unusually large: {mbap.length}")

    def _check_read_blocks(self, ctx: RequestContext, fc: int) -> None:
        tup = safe_unpack(">BHH", ctx.pdu)
        if not tup:
            return
        _, start, qty = tup
        if qty == 0:
            self._emit(ctx, "CVE-2023-25619", "异常条件检查不当，DoS", "qty=0 in read request")
        if start + qty > REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"read start+qty out of range: {start}+{qty}>{REG_CNT_MAX}")

    def _check_write_single_coil(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHH", ctx.pdu)
        if not tup:
            return
        _, addr, val = tup
        if addr >= REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"coil addr out of range: {addr}>{REG_CNT_MAX-1}")
        if val not in (0x0000, 0xFF00):
            self._emit(ctx, "CVE-2022-28613", "输入验证不当（重启）",
                       f"illegal coil value {val:#06x} (expected 0x0000 or 0xFF00)")

    def _check_write_single_reg(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHH", ctx.pdu)
        if not tup:
            return
        _, addr, _ = tup
        if addr >= REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"register addr out of range: {addr}>{REG_CNT_MAX-1}")

    def _check_write_multi_coils(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHHB", ctx.pdu)
        if not tup:
            return
        _, start, qty, bc = tup
        data = pdu_tail(ctx.pdu, 6)
        expected_bc = (qty + 7) // 8 if qty > 0 else 0
        if qty == 0:
            self._emit(ctx, "CVE-2023-25619", "异常条件检查不当，DoS", "qty=0 in 0x0F")
        if start + qty > REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"0x0F start+qty out of range: {start}+{qty}>{REG_CNT_MAX}")
        if bc > expected_bc:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x0F byte_count({bc}) > expected({expected_bc}) for qty={qty}")
        if len(data) < bc:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x0F declared bc({bc}) > available({len(data)})")

    def _check_write_multi_regs(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHHB", ctx.pdu)
        if not tup:
            return
        _, start, qty, bc = tup
        values = pdu_tail(ctx.pdu, 6)
        if qty == 0:
            self._emit(ctx, "CVE-2023-25619", "异常条件检查不当，DoS", "qty=0 in 0x10")
        if start + qty > REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"0x10 start+qty out of range: {start}+{qty}>{REG_CNT_MAX}")
        if bc > 2 * qty:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x10 byte_count({bc}) > 2*qty({2*qty})")
        if len(values) < bc:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x10 declared bc({bc}) > available({len(values)})")

    def _check_mask_write_reg(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHHH", ctx.pdu)
        if not tup:
            return
        _, addr, _, _ = tup
        if addr >= REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"0x16 addr out of range: {addr}>{REG_CNT_MAX-1}")

    def _check_read_write_multi(self, ctx: RequestContext) -> None:
        tup = safe_unpack(">BHHHHB", ctx.pdu)
        if not tup:
            return
        _, r_start, r_qty, w_start, w_qty, bc = tup
        values = pdu_tail(ctx.pdu, 11)
        if r_qty == 0 or w_qty == 0:
            self._emit(ctx, "CVE-2023-25619", "异常条件检查不当，DoS", f"r_qty={r_qty}, w_qty={w_qty}")
        if r_start + r_qty > REG_CNT_MAX or w_start + w_qty > REG_CNT_MAX:
            self._emit(ctx, "CVE-2024-34244", "缓冲区溢出（越界读）",
                       f"0x17 out of range: r:{r_start}+{r_qty} or w:{w_start}+{w_qty}>{REG_CNT_MAX}")
        if bc != 2 * w_qty:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x17 byte_count({bc}) != 2*w_qty({2*w_qty})")
        if len(values) < 2 * w_qty:
            self._emit(ctx, "CVE-2024-10918", "栈缓冲区溢出",
                       f"0x17 values too short: have {len(values)}, need {2*w_qty}")

# ----------------------------------
# Middleware wrapper
# ----------------------------------

def make_cve_middleware(probe: CVEProbe) -> Callable[[RequestContext], Awaitable[None]]:
    async def _mw(ctx: RequestContext) -> None:
        await probe.check(ctx)
    return _mw

# ----------------------------------
# Runner
# ----------------------------------

async def main_async() -> None:
    host = CONFIG["HOST"]
    port = _resolve_port()
    out_path = CONFIG["OUT"]
    ratelimit = int(CONFIG["RATELIMIT"])

    logger = CVELog(out_path)
    probe = CVEProbe(logger)
    middlewares = [mbs.mw_trace, mbs.rate_limiter(ratelimit), make_cve_middleware(probe)]
    server = mbs.ModbusTCPServer(host=host, port=port, datastore=mbs.DataStore(), middlewares=middlewares)

    await server.start()
    print(f"[{ts_hms()}] [CVE-Probe] Listening on {host}:{port}. JSONL: {out_path}. Ctrl+C to stop.")
    try:
        stop_ev = asyncio.Event()
        loop = asyncio.get_running_loop()
        for sig in (getattr(mbs.signal, "SIGINT", None), getattr(mbs.signal, "SIGTERM", None)):
            if sig is not None:
                try:
                    loop.add_signal_handler(sig, stop_ev.set)
                except NotImplementedError:
                    pass
        await stop_ev.wait()
    finally:
        await server.close()
        logger.close()
        if hasattr(probe, "counts"):
            totals = " | ".join(f"{k}:{probe.counts[k]}" for k in sorted(probe.counts))
            print(f"[{ts_hms()}] [CVE-Probe] Final totals -> {totals}")

def main() -> None:
    # 无命令行参数，直接使用 CONFIG
    try:
        asyncio.run(main_async())
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    main()
